from flask import Flask, render_template, jsonify, request, send_file
#from pdf2image import convert_from_path
import os, json, uuid, sqlite3
from datetime import datetime
from io import BytesIO
from PIL import Image
import fitz, pytesseract

app = Flask(__name__)

# 配置Tesseract路径（如果未添加到PATH）
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

@app.route("/")
def index():
    return render_template("index.html")

@app.route("/images_canvas")
def images_canvas():
    return render_template("images_canvas.html")

@app.route("/images_show")
def images_show():
    return render_template("images_show.html")

@app.route("/canvas_image_view")
def canvas_image_view():
    return render_template("CanvasImageView.html")

@app.route("/ollama")
def ollama():
    return render_template("ollama.html")

@app.route("/aichat")
def aichat():
    return render_template("aiChat.html")

@app.route("/ollama_ocr")
def ollama_ocr():
    return render_template("ollama_ocr.html")

@app.route("/test")
def test():
    return render_template("test.html")

def find_pdf_files(root_dir):
    """递归查找目录中的所有PDF文件"""
    pdf_files = {}
    for root, _, files in os.walk(root_dir):
        for file in files:
            if file.lower().endswith('.pdf'):
                full_path = os.path.join(root, file)
                doc = fitz.open(full_path)
                page_count = doc.page_count
                doc.close()
                # 统一使用正斜杠，方便JSON处理
                #normalized_path = full_path.replace('H:\\Books\\', '').replace('\\', '/')
                normalized_path = full_path.replace('H:\\Books\\', '')
                if normalized_path in pdf_files:
                    print(normalized_path)
                pdf_files[normalized_path] = page_count
                #pdf_files.append(normalized_path)
    return pdf_files
 
@app.route('/get_pdfs', methods=['GET'])
def get_pdfs():
    try:
        target_dir = r"H:\books\清华大学【图书馆书籍】\pdf_page_counts.json"
        with open(target_dir, "r", encoding="utf-8") as f:
            data = json.load(f)
            # 返回JSON响应
            return jsonify({
                'pdf_files': data
            })
    except FileNotFoundError:
        return jsonify({'error': f'❌ 错误：JSON 文件不存在！'}), 404
    except json.JSONDecodeError:
        return jsonify({'error': f'❌ 错误：JSON 文件格式无效！'}), 404

@app.route("/convert", methods=["POST"])
def convert_pdf_to_image():
    try:
        # 配置 PDF 文件存储路径（可选）
        UPLOAD_FOLDER = r"H:\books\清华大学【图书馆书籍】"
        os.makedirs(UPLOAD_FOLDER, exist_ok=True)
        # 1. 获取 POST 数据
        data = request.json
        filename = data.get("filename")
        page_num = int(data.get("page_number", 1)) - 1  # 转换为 0-based 索引
 
        # 2. 检查文件是否存在（这里假设 PDF 在服务器上）
        pdf_path = os.path.join(UPLOAD_FOLDER, filename)
        if not os.path.exists(pdf_path):
            return jsonify({"error": "File not found"}), 404
 
        # 3. 使用 PyMuPDF 读取 PDF 并转换指定页为图片
        doc = fitz.open(pdf_path)
        if page_num < 0 or page_num >= len(doc):
            return jsonify({"error": "Page number out of range"}), 400
 
        page = doc[page_num]
        pix = page.get_pixmap(dpi=150)  # 设置 DPI 控制质量
        img_path = os.path.join("static", f"page_{page_num}.png")
        pix.save(img_path)  # 保存图片（可选）
 
        # 4. 返回图片（这里直接返回文件，也可用 Base64 编码）
        return send_file(img_path, mimetype="image/png")
 
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 配置图片保存目录
UPLOAD_FOLDER = 'uploads'
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

@app.route('/ocr', methods=['POST'])
def ocr():
    if 'image' not in request.files:
        return jsonify({'error': '未找到图片文件'}), 400
 
    file = request.files['image']
    if file.filename == '':
        return jsonify({'error': '未选择文件'}), 400
    
    # 生成格式化的文件名：UUID + 时间戳 + 原始文件扩展名
    original_ext = os.path.splitext(file.filename)[1]
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    unique_id = str(uuid.uuid4().hex)[:8]  # 取UUID的前8位
    new_filename = f"{timestamp}_{unique_id}{original_ext}"
    # 保存图片
    save_path = os.path.join(UPLOAD_FOLDER, new_filename)

    try:
        img_bytes = file.read()
        img = Image.open(BytesIO(img_bytes))
        img.save(save_path)
        text = pytesseract.image_to_string(img, lang='chi_sim+eng')
        return jsonify({
            'filename': new_filename,
            'text': text.strip(),
            'success': True
        })
    except Exception as e:
        return jsonify({'error': f'OCR 处理失败: {str(e)}'}), 500

# 初始化数据库
def init_db():
    if not os.path.exists('第1章_数据库原理概述.db'):
        conn = sqlite3.connect('第1章_数据库原理概述.db')
        cursor = conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS ocr_ai (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                filename TEXT NOT NULL,
                model TEXT,
                ai TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        conn.commit()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS ocr_ai_messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                ocr_ai_id INTEGER NOT NULL,
                messages_id TEXT NOT NULL,
                messages_role TEXT NOT NULL,
                messages_content TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (ocr_ai_id) REFERENCES ocr_ai (id)
            )
        ''')
        conn.commit()
        conn.close()

@app.route('/upload_db', methods=['POST'])
def upload_data():
    if not request.is_json:
        return jsonify({"error": "Request must be JSON"}), 400
    
    data = request.get_json()
    
    try:
        # 验证数据结构
        
        ocr_ai_filename = data['filename']
        ocr_ai_model = data['model']
        ocr_ai_ai = data['ai']
        ocr_ai_messages = data['messages']
        
        # 插入OCR数据
        db = sqlite3.connect('第1章_数据库原理概述.db')
        cursor = db.cursor()
        
        cursor.execute(
            "INSERT INTO ocr_ai (filename, model, ai) VALUES (?, ?, ?)",
            (ocr_ai_filename, ocr_ai_model, ocr_ai_ai)
        )
        ocr_id = cursor.lastrowid
        
        db.commit()

        for index in range(len(ocr_ai_messages)):
            cursor.execute(
                "INSERT INTO ocr_ai_messages (ocr_ai_id, messages_id, messages_role, messages_content) VALUES (?, ?, ?, ?)",
                (ocr_id, index, ocr_ai_messages[index]["role"], ocr_ai_messages[index]["content"])
            )
        
        db.commit()
        
        return jsonify({
            "status": "success",
            "message": "Data saved successfully"
        }), 201
        
    except Exception as e:
        return jsonify({"error": str(e)}), 400
    finally:
        db.close()

if __name__ == '__main__':
    init_db()
    app.run(debug=True)